import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split, GridSearchCV
from sklearn.metrics import accuracy_score
import joblib

from . feature_engineering import feature_engineering_random_forest,load_config



def train_random_forest():
    X_train_red, X_test_red, y_train_red, y_test_red, X_train_blue, X_test_blue, y_train_blue, y_test_blue = feature_engineering_random_forest()
    
    config = load_config('ssq')

    # 定义超参数
    param_grid = {
        'n_estimators': config['model']['random_forest']['n_estimators'],  # 使用较大的树数量确保稳定性
        'max_depth': config['model']['random_forest']['max_depth'],
        'min_samples_split': config['model']['random_forest']['min_samples_split'],
        'min_samples_leaf': config['model']['random_forest']['min_samples_leaf']
    }

    #  多次训练并记录结果
    num_iterations = config['model']['random_forest']['num_iterations'] # 训练5次
    red_model_accuracies = []
    blue_model_accuracies = []

    for i in range(num_iterations):
        print(f'第 {i+1} 次训练')

        # 红球模型
        rf_model_red = RandomForestClassifier(random_state=42)
        grid_search_red = GridSearchCV(estimator=rf_model_red, param_grid=param_grid, cv=5, verbose=2, n_jobs=-1)
        grid_search_red.fit(X_train_red, y_train_red)
        best_model_red = grid_search_red.best_estimator_

        # 蓝球模型
        rf_model_blue = RandomForestClassifier(random_state=42)
        grid_search_blue = GridSearchCV(estimator=rf_model_blue, param_grid=param_grid, cv=5, verbose=2, n_jobs=-1)
        grid_search_blue.fit(X_train_blue, y_train_blue)
        best_model_blue = grid_search_blue.best_estimator_

        # 评估红球模型准确率
        y_pred_red = best_model_red.predict(X_test_red)
        accuracy_red = accuracy_score(y_test_red, y_pred_red)
        red_model_accuracies.append(accuracy_red)
        print(f'红球模型第 {i+1} 次准确率: {accuracy_red:.2f}')

        # 评估蓝球模型准确率
        y_pred_blue = best_model_blue.predict(X_test_blue)
        accuracy_blue = accuracy_score(y_test_blue, y_pred_blue)
        blue_model_accuracies.append(accuracy_blue)
        print(f'蓝球模型第 {i+1} 次准确率: {accuracy_blue:.2f}')

    # 打印多次训练的结果统计
    print(f'红球模型平均准确率: {np.mean(red_model_accuracies):.2f}, 方差: {np.var(red_model_accuracies):.2f}')
    print(f'蓝球模型平均准确率: {np.mean(blue_model_accuracies):.2f}, 方差: {np.var(blue_model_accuracies):.2f}')

    # 保存最后一次训练的模型
    joblib.dump(best_model_red, config['model']['random_forest']['model_red_path'])
    joblib.dump(best_model_blue, config['model']['random_forest']['model_blue_path'])